package usi.ch.onto.yago;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;

import usi.ch.onto.common.StringUtils;

public class ImportFacts {
	private String CONNECTION_STRING = "jdbc:mysql://localhost/yago";
	private int BATCH_MAX = 30000;
	private String USER = "root";
	private String PASS = "admin";
	private Connection conn;

	public ImportFacts() throws Exception {
		conn = DriverManager.getConnection(CONNECTION_STRING, USER, PASS);
		if (conn.isClosed())
			throw new Exception("Error connection to db ");

	}

	public void importTable(String tableDirectory) throws Exception {
		importTable(tableDirectory, false, false, false);
	}

	public void importTable(String tableDirectory, boolean parseArg1,
			boolean parseArg2, boolean parseAllFiles) throws Exception {
		importTable(tableDirectory, parseArg1, parseArg2, false, false,
				parseAllFiles, true, true);
	}

	public void importTable(String tableDirectory, boolean parseArg1,
			boolean parseArg2, boolean dontAddTerms, boolean useIsAExtractor,
			boolean parseAllFiles, boolean splitArg1, boolean splitArg2) throws Exception {
		File tableDir = new File(tableDirectory);
		// table name is the name of the directory
		if (tableDir.isDirectory() == false)
			throw new Exception(tableDirectory
					+ " does not point to a directory");
		String tableName = tableDir.getName();

		for (File fact : tableDir.listFiles()) {
			// each directory is a fact type
			if (fact.isDirectory() == false)
				continue;

			importFactType(fact, tableName, parseArg1, parseArg2, dontAddTerms,
					useIsAExtractor, parseAllFiles, splitArg1, splitArg2);
		}

	}

	private ExtractorType getExtractorTypeFromFileName(String filename) {
		ExtractorType facttype = null;
		filename = filename.toLowerCase();
		if (filename.contains("isaextractor"))
			facttype = ExtractorType.IsAExtractor;
		else if (filename.contains("articleextractor"))
			facttype = ExtractorType.ArticleExtractor;
		else if (filename.contains("checkedfactextractor"))
			facttype = ExtractorType.CheckedFactExtractor;
		else if (filename.contains("wordnetlinks"))
			facttype = ExtractorType.WordNetLinks;
		else if (filename.contains("conceptlinker"))
			facttype = ExtractorType.ConceptLinker;

		return facttype;
	}

	private void importFactType(File fact, String tablename, boolean parseArg1,
			boolean parseArg2, boolean dontAddTerms, boolean useIsAExtractor,
			boolean parseAllFiles, boolean splitArg1, boolean splitArg2) throws Exception {
		ExtractorType facttype = null;
		File fileToImport = null;
		for (File article : fact.listFiles()) {
			if (parseAllFiles) {
				facttype = getExtractorTypeFromFileName(article.getName());
				if (facttype != null) {
					importFile(article, tablename, fact.getName(), facttype,
							parseArg1, parseArg2, dontAddTerms, splitArg1, splitArg2);
				}
			} else if (useIsAExtractor) {
				if (facttype == ExtractorType.IsAExtractor) {
					fileToImport = article;
					break;
				}
			} else {
				if (facttype == ExtractorType.ArticleExtractor) {
					fileToImport = article;
					break; // priority file
				} else if (facttype == ExtractorType.CheckedFactExtractor)
					fileToImport = article; // not priority
			}
		}
		if (! parseAllFiles && fileToImport != null)
			importFile(fileToImport, tablename, fact.getName(), facttype,
					parseArg1, parseArg2, dontAddTerms, splitArg1, splitArg2);

	}

	private void importFile(File file, String table, String factName,
			ExtractorType facttype, boolean parseArg1, boolean parseArg2,
			boolean dontAddTerms, boolean splitArg1, boolean splitArg2) throws Exception {
		BufferedReader br = new BufferedReader(new FileReader(file));
		String line = null;
		String[] fields;

		try {
			PreparedStatement pstmtTable = conn
					.prepareStatement(getTableInsertString(table));
			PreparedStatement pstmtTermsArg1 = null;
			PreparedStatement pstmtTermsArg2 = null;
			if (dontAddTerms == false) {
				pstmtTermsArg1 = conn
						.prepareStatement(getTableArg1TermsInsertString(table));
				pstmtTermsArg2 = conn
						.prepareStatement(getTableArg2TermsInsertString(table));
			}
			conn.setAutoCommit(false);
			int batchCount = 0;
			while ((line = br.readLine()) != null) {
				fields = line.split("\t");

				if (fields.length != 4) {
					System.out.println("invalid line " + line);
					continue;
				}
				insertSingleFact(pstmtTable, pstmtTermsArg1, pstmtTermsArg2,
						fields, table, factName, facttype, parseArg1,
						parseArg2, dontAddTerms, splitArg1, splitArg2);
				batchCount++;

				if (batchCount >= BATCH_MAX) {
					batchCount = 0;
					executeInsertBatch(pstmtTable, pstmtTermsArg1,
							pstmtTermsArg2, dontAddTerms);
				}
			}
			executeInsertBatch(pstmtTable, pstmtTermsArg1, pstmtTermsArg2,
					dontAddTerms);
			pstmtTable.close();

		} catch (IOException iox) {
			throw iox;
		}
	}

	// this is batch for model 2 (Multiple term in tables %table%Arg1Terms and
	// %table%Arg2Terms
	private void executeInsertBatch(PreparedStatement pstmtTable,
			PreparedStatement pstmtTermsArg1, PreparedStatement pstmtTermsArg2,
			boolean dontAddTerms) throws Exception {
		pstmtTable.executeBatch();
		if (dontAddTerms == false) {
			pstmtTermsArg1.executeBatch();
			pstmtTermsArg2.executeBatch();
		}
		conn.commit();
		System.out.println("batch completed of " + BATCH_MAX);
	}

	private String getTableInsertString(String tablename) throws Exception {
		return "insert into "
				+ tablename
				+ "(id, arg1, arg2, factName, extractorType, confidence, isArg1SingleTerm, isArg2SingleTerm) values(?, ?, ?, ?, ?, ?, ?, ?);";
	}

	private String getTableArg1TermsInsertString(String tablename)
			throws Exception {
		return "insert into " + tablename + "Arg1Terms"
				+ "(tabid, Term) values(?, ?);";
	}

	private String getTableArg2TermsInsertString(String tablename)
			throws Exception {
		return "insert into " + tablename + "Arg2Terms"
				+ "(tabid, Term) values(?, ?);";
	}

	// inserts to
	private void insertSingleFact(PreparedStatement pstmt,
			PreparedStatement pstmtArg1, PreparedStatement pstmtArg2,
			String[] fields, String table, String factName,
			ExtractorType extracttype, boolean parseArg1, boolean parseArg2,
			boolean dontAddTerms, boolean splitArg1, boolean splitArg2) throws Exception {
		long singleFactID = Long.parseLong(fields[0]);
		String arg1 = fields[1];
		String arg2 = fields[2];
		short extractorValue = FactExtractor.getExtractorTypeValue(extracttype);
		float confidence = Float.parseFloat(fields[3]);

		// insert into %table%, %table%Terms, %table%_FactTerms
		pstmt.setLong(1, singleFactID);
		pstmt.setString(2, arg1);
		pstmt.setString(3, arg2);
		pstmt.setString(4, factName);
		pstmt.setShort(5, extractorValue);
		pstmt.setFloat(6, confidence);
		if (dontAddTerms){
			//if we dont add terms we consider that they are treated as a single term
			pstmt.setBoolean(7, true); //isArg1..
			pstmt.setBoolean(8, true); //isArg2..
		}else {
			List<String> list = getTermsFromString(arg1, parseArg1, splitArg1);
			addTermsBatch(pstmtArg1, list, singleFactID);
			pstmt.setBoolean(7, list.size() == 1); //isArg1SingleTerm
						
			list = getTermsFromString(arg2, parseArg2, splitArg2);
			addTermsBatch(pstmtArg2, list, singleFactID);
			pstmt.setBoolean(8, list.size() == 1); //isArg2SingleTerm
		}
		pstmt.addBatch();

	}

	private void addTermsBatch(PreparedStatement stmt, List<String> listTerms,
			long singleFactID) throws SQLException {
		for (String term : listTerms) {
			// only add non empty term
			if (term.length() > 0) {
				stmt.setLong(1, singleFactID);
				stmt.setString(2, term);
				stmt.addBatch();
			}
		}
	}

	// returns array of single terms
	// splitted from _ removed also "
	private List<String> getTermsFromString(String sentence, boolean parseTerms, boolean splitTerms) {
		List<String> terms;
		String strParsed;
		if (parseTerms)
			strParsed = parseSingleTerm(sentence);
		else
			strParsed = sentence;
		
		if (splitTerms) {
			terms = StringUtils.fastSplit(
					strParsed.replace(' ', '_').replace("\"", ""), '_', ' ');			
		} else {
			terms = new ArrayList<String>();
			terms.add(parseSingleTerm(strParsed));
		}

		return terms;
	}

	private String parseSingleTerm(String sentence) {
		String ret = "";
		if (sentence.startsWith("wikicategory")) {
			// this ignores already the _ char
			ret = sentence.substring("wikicategory".length() + 1);
		} else if (sentence.startsWith("wordnet")) {
			ret = sentence.substring("wordnet".length() + 1,
					sentence.lastIndexOf("_"));
		}

		return ret;
	}

	// right now just check for the strings wikicategory wordnet
	// wikicategory is removed
	// wordnet and last string is also removed
	// they are considered to be the first terms
	private String[] parseTermsGetMultipleTerms(String[] initialTermList) {
		if (initialTermList.length > 1) {
			if (initialTermList[0].equals("wikicategory")) {
				initialTermList[0] = "";
			} else if (initialTermList[0].equals("wordnet")) {
				initialTermList[0] = "";
				initialTermList[initialTermList.length - 1] = "";
			}
		}
		return initialTermList;
	}

}